package my.proxy;

import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class TcpProxyServer {
    private static final Logger log = LoggerFactory.getLogger(TcpProxyServer.class);

    private final int port;

    public TcpProxyServer(int port) {
        this.port = port;
    }

    public static void main(String[] args) {
        new TcpProxyServer(8888).run();
    }

    public void run() {
        log.info("--------TCP Proxy start-----------");
        EventLoopGroup bossGroup = new NioEventLoopGroup(1);
        EventLoopGroup workerGroup = new NioEventLoopGroup();
        LoggingHandler debugHandler = new LoggingHandler(LogLevel.DEBUG);

        Bootstrap remoteBootstrap = new Bootstrap();
        remoteBootstrap.group(workerGroup)
                .channel(NioSocketChannel.class)
                .option(ChannelOption.AUTO_READ, false)
                .option(ChannelOption.TCP_NODELAY, true)
                .remoteAddress("10.8.0.6", 8888)
                .handler(debugHandler);

        try {
            ServerBootstrap b = new ServerBootstrap();
            b.group(bossGroup, workerGroup)
                    .channel(NioServerSocketChannel.class)
                    .option(ChannelOption.TCP_NODELAY, true)
                    .handler(new LoggingHandler(LogLevel.INFO))
                    .childHandler(new ChannelInitializer<SocketChannel>() {
                        @Override
                        public void initChannel(SocketChannel ch) {
                            ch.pipeline().addLast(
                                    debugHandler,
                                    new TcpProxyClientHandler(remoteBootstrap));
                        }
                    }).bind(port).sync().channel().closeFuture().sync();
        } catch (Throwable e) {
            log.error(e.getMessage(), e);
        } finally {
            bossGroup.shutdownGracefully();
            workerGroup.shutdownGracefully();
        }
    }

    /**
     * 代理客户端去请求目标主机
     */
    private static class TcpProxyClientHandler extends ChannelInboundHandlerAdapter {

        private Channel remoteChannel;
        private final Bootstrap remoteBootstrap;

        public TcpProxyClientHandler(Bootstrap remoteBootstrap) {
            this.remoteBootstrap = remoteBootstrap;
        }

        @Override
        public void channelActive(ChannelHandlerContext ctx) {
            Channel clientChannel = ctx.channel();
            clientChannel.config().setAutoRead(false);
            ChannelFuture f = remoteBootstrap.connect();
            remoteChannel = f.channel();
            remoteChannel.pipeline().addLast(new TcpProxyRemoteHandler(clientChannel));
            remoteChannel.config().setAutoRead(true);
            f.addListener((ChannelFutureListener) future -> {
                if (future.isSuccess()) {
                    clientChannel.config().setAutoRead(true);
                } else {
                    clientChannel.close();
                }
            });
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) {
            remoteChannel.writeAndFlush(msg);
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) {
            flushAndClose(remoteChannel);
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable e) {
            log.error(e.getMessage(), e);
            flushAndClose(ctx.channel());
            flushAndClose(remoteChannel);
        }

        private void flushAndClose(Channel ch) {
            if (ch != null && ch.isActive()) {
                ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            }
        }
    }

    /**
     * 代理客户端请求目标主机处理器
     */
    private static class TcpProxyRemoteHandler extends ChannelInboundHandlerAdapter {

        private final Channel clientChannel;

        public TcpProxyRemoteHandler(Channel clientChannel) {
            this.clientChannel = clientChannel;
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) {
            clientChannel.writeAndFlush(msg);
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) {
            flushAndClose(clientChannel);
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable e) {
            log.error(e.getMessage(), e);
            flushAndClose(ctx.channel());
            flushAndClose(clientChannel);
        }

        private void flushAndClose(Channel ch) {
            if (ch != null && ch.isActive()) {
                ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            }
        }
    }
}
